import pandas as pd
import subprocess
import sys
import os
import argparse
import time  # 导入time模块用于暂停程序
from subprocess import TimeoutExpired  # 导入超时异常类

def run(args):
    C = args.C
    device_id = args.D
    # sigtype,m,n DataFrame Columns
    operator_path = args.performance_file
    data_path = args.data_path
    data = pd.read_csv('./params/{}'.format(data_path))
    results = pd.DataFrame(columns=["sigtype", "m", "n", "time_us"])
    prof_data_path = "./batch_prof_data.csv"
    for index, row in data.iterrows():
        if index > C:
            break
        sigtype = row.iloc[0]
        m = row.iloc[1]
        n = row.iloc[2]
        command = 'rm -rf ./prof/* && mkdir -p ./prof'
        result = subprocess.run(command, shell=True)
        # warmup 能跑到峰值性能就行 现在这个配置是能跑到1800的
        # msprof op --kernel-name=softmax_kernel pytest performance_softmax.py
        shape = str(m) + " " + str(n)
        # command = "python {} --dtype {} --dims {} --seed {} --device {}".format(operator_path, sigtype, shape, args.seed, device_id)
        command = "python {} --dtype {} --dims {} --seed {} --device {}".format(operator_path, sigtype, shape, args.seed, device_id)
        print(f"[用例{index}] 运行命令：{command}")
        try:
            result = subprocess.run(command, shell=True, capture_output=True, text=True, timeout=60, check=False)
        except TimeoutExpired:
            command = 'python3 ./prof.py $(find ./prof -name "op_statistic.csv") {} {} {}'.format(sigtype, m, n)
            result = subprocess.run(command, shell=True, capture_output=True, text=True)
            last_line1 = result.stdout.strip().splitlines()[-1]
            parts = last_line1.split()  
            sigtype = parts[1]
            m = parts[3]
            n = parts[5]
            time_us = parts[7]
            frame =  pd.DataFrame([[sigtype, m, n, time_us]], 
                                columns=["sigtype", "m", "n", "time_us"])

            if index == 0:
                frame.to_csv(prof_data_path, mode='w', header=True, index=False)
            else:
                frame.to_csv(prof_data_path, mode='a', header=not os.path.exists(prof_data_path), index=False)
            time.sleep(1)  # 在每次循环结束后暂停1秒
            # current_status = "test_timeout"
            # pass
            # print(f"[用例{index}] 测试命令超时（{test_timeout}秒），终止运行：{sigtype} {m}x{n}")
        # result = subprocess.run(command, shell=True, capture_output=True, text=True, timeout=60, check=False)
        # command = 'python3 ./prof.py $(find ./prof -name "OpBasicInfo*.csv") $(find ./prof -name "PipeUtilization*.csv") {} {} {}'.format(sigtype, m, n)
        command = 'python3 ./prof.py $(find ./prof -name "op_statistic.csv") {} {} {}'.format(sigtype, m, n)
        result = subprocess.run(command, shell=True, capture_output=True, text=True)
        last_line1 = result.stdout.strip().splitlines()[-1]
        print(last_line1)
        parts = last_line1.split()  
        sigtype = parts[1]
        m = parts[3]
        n = parts[5]
        time_us = parts[7]
        frame =  pd.DataFrame([[sigtype, m, n, time_us]], 
                            columns=["sigtype", "m", "n", "time_us"])

        if index == 0:
            frame.to_csv(prof_data_path, mode='w', header=True, index=False)
        else:
            frame.to_csv(prof_data_path, mode='a', header=not os.path.exists(prof_data_path), index=False)
        time.sleep(1)  # 在每次循环结束后暂停1秒
        

    
if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Silu Prof")
    parser.add_argument("--performance_file", type=str, default="test_silu.py", help="The file to save the prof data")
    parser.add_argument("--data_path", type=str, default="silu_performance_data.csv", help="The path of the test data")
    parser.add_argument("--C", type=int, default=100, help="The length of the test data")
    parser.add_argument("--D", type=int, default=7, help="Device ID")
    parser.add_argument("--seed", type=int, default=42, help="Random seed")
    args = parser.parse_args()
    run(args)
